from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())    # read local .env file

import streamlit as st
from langchain.prompts import PromptTemplate
from langchain_core.messages import AIMessage, HumanMessage
from chat.chat import GetAgentExecutor





# Streamlit 应用程序界面
def main():
    # 主题设为light
    st.set_page_config(
        layout="wide", 
        page_title="🤵‍♀️个人理财助理📟", 
        page_icon="🤖",
        initial_sidebar_state="expanded",

        )
    st.title('🤵‍♀️个人理财助理📟')
    Zhipu_api_key = st.sidebar.text_input('智谱 API Key', type='password')
    Tavily_api_key = st.sidebar.text_input('Tavily API Key', type='password')

    # Siderbar上一个确认是否使用上下文记忆的选项
    use_memory = st.sidebar.checkbox("使用上下文记忆", value=True)
    # Siderbar上一个确认是否使用知识库
    use_document = st.sidebar.checkbox("使用知识库", value=True)
    # Siderbar三个可选模型，分别为GLM-4-0520、ChatGLM-3-6B微调和GLM-3-Turbo微调,返回选中下标
    selected_method = st.sidebar.selectbox(
        "选择模型",
        ["GLM-4-0520", "ChatGLM-3-6B微调", "GLM-3-Turbo微调"],
    )

    # Sidebar上一个按钮，用于更新

    selected_index = {
        "GLM-4-0520": 0,
        "ChatGLM-3-6B微调": 1,
        "GLM-3-Turbo微调": 2,
    }[selected_method]

    # 用于跟踪对话历史
    if 'messages' not in st.session_state:
        st.session_state.messages = []

    messages = st.container(border=True)
    if prompt := st.chat_input("Say something"):
        # 将用户输入添加到对话历史中
        st.session_state.messages.append({"role": "user", "text": prompt})
        chat_history = []
        if use_memory:
            # 将最后三个用户消息添加到对话历史中
            length = len(st.session_state.messages)
            mk = min((length - 1) // 2, 3)
            for i in range(1, mk + 1):
                chat_history.extend(
                    [
                        HumanMessage(content=st.session_state.messages[-2 * i - 1]["text"]),
                        AIMessage(content=st.session_state.messages[-2 * i]["text"]),
                    ]
                )
        answer = GetAgentExecutor(
            option=selected_index
        ).invoke({"input": prompt, "chat_history": chat_history, "use_document": use_document})
        print(answer["output"])

        # 检查回答是否为 None
        if answer is not None:
            # 将LLM的回答添加到对话历史中
            st.session_state.messages.append({"role": "assistant", "text": answer['output']})

        # 显示整个对话历史
        for message in st.session_state.messages:
            if message["role"] == "user":
                messages.chat_message("user").write(message["text"])
            elif message["role"] == "assistant":
                messages.chat_message("assistant").write(message["text"])   


if __name__ == "__main__":
    main()
